package com.yjzx.util.common.util.file;

import com.yjzx.util.common.util.ExceptionUtil;
import com.yjzx.util.common.util.map.ListNode;
import com.yjzx.util.common.util.map.NodeMap;
import com.yjzx.util.common.util.exception.BiConsumerThrow;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

/**
 * @author yjzx
 * @date 2024/6/4
 * @description ZipUtils
 */
public class ZipUtils {
    /**
     * zip文件追加
     *
     * @param outputStream 输出流，要写入的地方
     * @param inputStream  输入流，要读取的zip流
     * @param filesBytes   要写入zip的文件
     * @throws IOException
     */
    public static void zipAddToInputStream(OutputStream outputStream, InputStream inputStream, ListNode<String, InputStream> filesBytes) throws IOException {
        try (ZipOutputStream zipOutputStream = outputStream instanceof ZipOutputStream ? (ZipOutputStream) outputStream : new ZipOutputStream(outputStream)) {
            try (ZipInputStream zipInputStream = inputStream instanceof ZipInputStream ? (ZipInputStream) inputStream : new ZipInputStream(inputStream)) {
                ZipEntry entry;
                while ((entry = zipInputStream.getNextEntry()) != null) {
                    if (!entry.isDirectory()) {
                        // 读取文件内容并存储
                        byte[] content = FileUtils.readFully(zipInputStream);
                        packageToZip(zipOutputStream, new ListNode<>(entry.getName(), content));
                    }
                    zipInputStream.closeEntry();
                }
            }
            packageFileToZip(zipOutputStream, filesBytes);
        }
    }

    /**
     * 同名添加文件会修改
     *
     * @param outputStream 输出流，要写入的地方
     * @param inputStream  输入流，要读取的zip流
     * @param filesBytes   要写入zip的文件同名会修改
     * @throws IOException
     */
    public static void zipUpdateToInputStream(OutputStream outputStream, InputStream inputStream, ListNode<String, byte[]> filesBytes) throws IOException {
        try (ZipOutputStream zipOutputStream = outputStream instanceof ZipOutputStream ? (ZipOutputStream) outputStream : new ZipOutputStream(outputStream)) {
            try (ZipInputStream zipInputStream = inputStream instanceof ZipInputStream ? (ZipInputStream) inputStream : new ZipInputStream(inputStream)) {
                Set<String> updateFileNameSet = filesBytes.stream().map(NodeMap::getKey).collect(Collectors.toSet());
                ZipEntry entry;
                while ((entry = zipInputStream.getNextEntry()) != null) {
                    if (!entry.isDirectory()) {
                        // 读取文件内容并存储
                        byte[] content = FileUtils.readFully(zipInputStream);
                        String name = entry.getName();
                        if (!updateFileNameSet.contains(name)) {
                            packageToZip(zipOutputStream, new ListNode<>(entry.getName(), content));
                        }
                    }
                    zipInputStream.closeEntry();
                }
            }
            packageToZip(zipOutputStream, filesBytes);
        }
    }

    /**
     * 输入流写入zip
     *
     * @param outputStream 输出流，要写入的地方
     * @param files        要写入的文件
     * @param function     要写入zip的文件
     * @throws IOException
     */
    public static <T, E extends Exception> void packageToZip(OutputStream outputStream, ListNode<String, T> files, BiConsumerThrow<ZipOutputStream, T, E> function) throws IOException, E {
        if (outputStream != null) {
            if (!(outputStream instanceof ZipOutputStream)) {
                try (ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream)) {
                    toZip(zipOutputStream, files, function);
                }
            } else {
                toZip((ZipOutputStream) outputStream, files, function);
            }
        }
    }

    private static <T, E extends Exception> void toZip(ZipOutputStream zipOutputStream, ListNode<String, T> files, BiConsumerThrow<ZipOutputStream, T, E> function) throws IOException, E {
        for (NodeMap<String, T> file : files) {
            String fileName = file.getKey();
            T t = file.getValue();
            //t为空时，跳过此文件
            if (fileName != null && t != null) {
                // 依次对每个文件都压缩
                zipOutputStream.putNextEntry(new ZipEntry(fileName));
                function.accept(zipOutputStream, t);
                zipOutputStream.closeEntry();
            }
        }
    }

    /**
     * 压缩多个byte数据文件写入zip
     *
     * @param outputStream
     * @param files
     * @throws IOException
     */
    public static void packageToZip(OutputStream outputStream, ListNode<String, byte[]> files) throws IOException {
        packageToZip(outputStream, files, FilterOutputStream::write);
    }

    /**
     * 压缩多个bufferedImage为zip包文件流
     *
     * @param files 文件名称 文件流
     * @return
     */
    public static void packageImageToZip(OutputStream outputStream, ListNode<String, BufferedImage> files) throws IOException {
        packageToZip(outputStream, files, (x, y) -> ImageIO.write(y, "png", x));
    }


    /**
     * 压缩多个文件为zip包文件流
     *
     * @param outputStream zip压缩文件的输出流
     * @param files        文件名称 文件输入流
     * @return
     */
    public static void packageFileToZip(OutputStream outputStream, ListNode<String, InputStream> files) throws IOException {
        byte[] buffer = new byte[0X400];
        packageToZip(outputStream, files, (x, y) -> {
            int len;
            while ((len = y.read(buffer)) > 0) {
                x.write(buffer, 0, len);
            }
        });
    }


    public static void throwException(Exception exception) {
        ExceptionUtil.throwToRunException(exception, RuntimeException.class);
    }
}
